use core::alloc::{GlobalAlloc, Layout};

use syscall::num::SyscallIndex;

#[naked]
extern "C" fn malloc_syscall(
    _id: u64,
    _arg1: usize,
    _arg2: usize,
    _arg3: usize,
    _arg4: usize,
    _arg5: usize,
) -> usize {
    unsafe {
        core::arch::naked_asm!(
            "mov rax, rdi",
            "mov rdi, rsi",
            "mov rsi, rdx",
            "mov rdx, rcx",
            "mov r10, r8",
            "mov r8, r9",
            "syscall",
            "ret",
        )
    }
}

fn malloc(layout: Layout) -> Result<u64, ()> {
    let addr = malloc_syscall(
        SyscallIndex::Malloc as u64,
        layout.size(),
        layout.align(),
        0,
        0,
        0,
    );

    if addr == 0 {
        Err(())
    } else {
        Ok(addr as u64)
    }
}

fn free(addr: u64, layout: Layout) {
    malloc_syscall(
        SyscallIndex::Free as u64,
        addr as usize,
        layout.size(),
        layout.align(),
        0,
        0,
    );
}

struct MemoryAllocator;

unsafe impl GlobalAlloc for MemoryAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        malloc(layout).unwrap() as *mut u8
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        free(ptr as u64, layout)
    }
}

#[global_allocator]
static GLOBAL_ALLOCATOR: MemoryAllocator = MemoryAllocator;
